#!/usr/bin/env python3
"""
快速下载Bark模型
"""

import os
import sys

# 设置环境变量
os.environ['XDG_CACHE_HOME'] = os.path.expanduser('~/.cache')
os.environ['XDG_DATA_HOME'] = os.path.expanduser('~/.local/share')
os.environ['TTS_CACHE_DIR'] = os.path.expanduser('~/.cache/tts')
os.environ['TORCH_HOME'] = os.path.expanduser('~/.cache/torch')
os.environ['HF_HOME'] = os.path.expanduser('~/.cache/huggingface')

# 创建目录
os.makedirs(os.path.expanduser('~/.cache/tts'), exist_ok=True)

try:
    from TTS.api import TTS
    print("✅ TTS库导入成功")
except ImportError as e:
    print(f"❌ TTS库导入失败: {e}")
    sys.exit(1)

def download_bark():
    """下载Bark模型"""
    print("🔄 开始下载Bark模型...")
    print("📍 这是多人对话的核心模型，可能需要几分钟...")
    
    try:
        # 下载Bark模型
        tts = TTS(model_name="tts_models/multilingual/multi-dataset/bark", progress_bar=True)
        print("✅ Bark模型下载成功！")
        
        # 测试模型
        print("🧪 测试模型...")
        test_text = "你好，这是一个测试。"
        tts.tts_to_file(
            text=test_text,
            file_path="/tmp/test_bark.wav",
            voice_preset="v2/zh_speaker_7"
        )
        print("✅ 模型测试成功！")
        
        # 清理测试文件
        if os.path.exists("/tmp/test_bark.wav"):
            os.remove("/tmp/test_bark.wav")
            
        return True
        
    except Exception as e:
        print(f"❌ Bark模型下载失败: {e}")
        return False

if __name__ == "__main__":
    print("="*50)
    print("🎭 Bark模型下载器")
    print("="*50)
    
    if download_bark():
        print("\n🎉 Bark模型准备完成！")
        print("现在可以测试多人对话功能了。")
    else:
        print("\n💥 下载失败，请检查网络连接")
        sys.exit(1)